///
/// Copyright (C) 2016, Dependable Systems Laboratory, EPFL
/// Copyright (C) 2015-2016, Cyberhaven
/// All rights reserved.
///
/// Licensed under the Cyberhaven Research License Agreement.
///

#include <s2e/cpu.h>

extern "C" {
#include "qdict.h"
#include "qint.h"
#include "qjson.h"
#include "qlist.h"
}

#include <klee/UserSearcher.h>
#include <llvm/Support/FileSystem.h>
#include <llvm/Support/Path.h>
#include <klee/util/ExprUtil.h>
#include <s2e/ConfigFile.h>
#include <s2e/S2E.h>
#include <s2e/S2EExecutor.h>

#include <random>
#include <sstream>
#include <system_error>

#include <boost/regex.hpp>
#include <capstone/capstone.h>

#include "GenerationSearcher.h"

namespace s2e {
namespace plugins {
namespace generations {

S2E_DEFINE_PLUGIN(GenerationSearcher, "Generation searcher", "", "MultiSearcher", "CUPASearcher");

namespace {

class GenerationSearcherState : public PluginState {
public:

    GenerationSearcherState() {};

    virtual GenerationSearcherState *clone() const {
        GenerationSearcherState *ret = new GenerationSearcherState(*this);
        return ret;
    }

    static PluginState *factory(Plugin *p, S2EExecutionState *s) {
        return new GenerationSearcherState();
    }

    virtual ~GenerationSearcherState() {
    }
};
}

GenerationSearcher::~GenerationSearcher(){

}
void GenerationSearcher::initialize() {
	m_detector = s2e()->getPlugin<ModuleExecutionDetector>();
	m_tc = s2e()->getPlugin<testcases::TestCaseGenerator>();
    ConfigFile *cfg = s2e()->getConfig();
   	ConfigFile::string_list pollingEntries = cfg->getListKeys(
   			getConfigKey() + ".forkRanges");
	Range rg;
	rg.start = 0x00000000;
	rg.end = 0x00000000;
	m_forkRanges.insert(rg);
   	foreach2(it, pollingEntries.begin(), pollingEntries.end())
   	{
   		std::stringstream ss1;
   		ss1 << getConfigKey() << ".forkRanges" << "." << *it;
   		ConfigFile::integer_list il = cfg->getIntegerList(ss1.str());
   		if (il.size() != 2) {
   			s2e()->getDebugStream() << "Range entry " << ss1.str()
   					<< " must be of the form {startPc, endPc} format" << '\n';
   			continue;
   		}

   		bool ok = false;
   		uint64_t start = cfg->getInt(ss1.str() + "[1]", 0, &ok);
   		if (!ok) {
   			s2e()->getDebugStream() << "ForkLimiter could not read "
   					<< ss1.str() << "[0]" << '\n';
   			continue;
   		}

   		uint64_t end = cfg->getInt(ss1.str() + "[2]", 0, &ok);
   		if (!ok) {
   			s2e()->getDebugStream() << "ForkLimiter could not read "
   					<< ss1.str() << "[1]" << '\n';
   			continue;
   		}
   		Range rg;
   		rg.start = start;
   		rg.end = end;
   		m_forkRanges.insert(rg);
   	}

    ConfigFile::integer_list forkHistory = cfg->getIntegerList(getConfigKey() + ".forkHistory");
	foreach2(it, forkHistory.begin(), forkHistory.end())
   	{
		m_forkHistory.push_back(*it);
   	}
    s2e()->getDebugStream() << "ForkHistory Size="
       					<< hexval(m_forkHistory.size())<< '\n';

    ConfigFile::integer_list onlyForkConditions = cfg->getIntegerList(getConfigKey() + ".onlyForkConditions");
	foreach2(it, onlyForkConditions.begin(), onlyForkConditions.end())
	{
		m_onlyForkConditions.insert(*it);
	}
	s2e()->getDebugStream() << "OnlyForkConditions Size="
						<< hexval(m_onlyForkConditions.size())<< '\n';

    ConfigFile::integer_list disableForkConditions = cfg->getIntegerList(getConfigKey() + ".disableForkConditions");
	foreach2(it, disableForkConditions.begin(), disableForkConditions.end())
	{
		m_disableForkConditions.insert(*it);
	}
	s2e()->getDebugStream() << "DisableForkConditions Size="
						<< hexval(m_disableForkConditions.size())<< '\n';

	ConfigFile::integer_list onlyForkPoints = cfg->getIntegerList(getConfigKey() + ".onlyForkPoints");
	foreach2(it, onlyForkPoints.begin(), onlyForkPoints.end())
	{
		m_configedOnlyForkPoints.insert(*it);
	}
	s2e()->getDebugStream() << "OnlyForkPoints Size="
						<< hexval(m_configedOnlyForkPoints.size())<< '\n';


    CorePlugin *plg = s2e()->getCorePlugin();
    plg->onStateFork.connect(sigc::mem_fun(*this, &GenerationSearcher::onStateFork));
    plg->onStateKill.connect(sigc::mem_fun(*this, &GenerationSearcher::onStateKill));
    plg->onSymbolicAddress.connect(sigc::mem_fun(*this, &GenerationSearcher::onSymbolicAddress));
    if (m_detector) {
		s2e()->getCorePlugin()->onStateForkDecide.connect(sigc::mem_fun(*this, &GenerationSearcher::onStateForkDecide));
		m_detector->onModuleTranslateBlockStart.connect(sigc::mem_fun(*this, &GenerationSearcher::onModuleTranslateBlockStart));
	    m_detector->onModuleTranslateBlockEnd.connect(
	        sigc::mem_fun(*this, &GenerationSearcher::onModuleTranslateBlockEnd));
	}
    s2e()->getCorePlugin()->onUnableFork.connect(sigc::mem_fun(*this, &GenerationSearcher::onUnableFork));

    m_underConstraint = cfg->getBool(getConfigKey() + ".underConstraint", false);
    m_searchBug = cfg->getBool(getConfigKey() + ".searchBug", false);
    m_genMultiHitCase = cfg->getBool(getConfigKey() + ".genMultiHitCase", false);
    m_maxUnForkGen = s2e()->getConfig()->getInt(getConfigKey() + ".maxUnForkGen", 10);

    m_cupa = s2e()->getPlugin<CUPASearcher>();
    m_multiSearcher = s2e()->getPlugin<MultiSearcher>();
    m_multiSearcher->registerSearcher("GenerationSearcher", this);
    switchToGenerationSearcher();
    m_initialState = nullptr;
}

void GenerationSearcher::onModuleTranslateBlockStart(ExecutionSignal *signal, S2EExecutionState *state,
                                               const ModuleDescriptor &module, TranslationBlock *tb, uint64_t pc) {
	if(m_onlyForkConditions.size() == 0 ) return;

    uint64_t native_pc_bb_start;
    if (!module.ToNativeBaseRebased(pc, native_pc_bb_start)) {
        getWarningsStream(state) << "Could not convert " << hexval(native_pc_bb_start) << " to relative address for module "
                                 << module.Name << "\n";
        return;
    }
    //only for bb_start
    if (m_onlyForkConditions.count(native_pc_bb_start)) {
    	getDebugStream(state) << "Found onlyForkConditions for pc " << hexval(native_pc_bb_start) << "\n";
    	signal->connect(sigc::mem_fun(*this, &GenerationSearcher::onOnlyForkTargetExecution));
    }

}

void GenerationSearcher::onModuleTranslateBlockEnd(ExecutionSignal *signal, S2EExecutionState *state,
                                                     const ModuleDescriptor &module, TranslationBlock *tb,
                                                     uint64_t endPc, bool staticTarget, uint64_t targetPc) {
	if(m_onlyForkConditions.size() == 0 ) return;

    uint64_t pc;
    if (!module.ToNativeBaseRebased(endPc, pc)) {
        getWarningsStream(state) << "Could not convert " << hexval(endPc) << " to relative address for module "
                                 << module.Name << "\n";
        return;
    }
    //only for bb_end
    if (m_onlyForkConditions.count(pc)) {
    	getDebugStream(state) << "Found onlyForkConditions for pc " << hexval(pc) << "\n";
    	signal->connect(sigc::mem_fun(*this, &GenerationSearcher::onOnlyForkTargetExecution));
    }

    //both for bb_end or bb_sart
	if (m_disableForkConditions.count(pc)) {
		getDebugStream(state) << "Found disableForkConditions for pc "
				<< hexval(pc) << "\n";
		signal->connect(
				sigc::mem_fun(*this,
						&GenerationSearcher::onDisableForkTargetExecution));
	}
    uint64_t bb_start_pc;
    if (module.ToNativeBaseRebased(tb->pc, bb_start_pc)) {
		if (m_disableForkConditions.count(bb_start_pc)) {
			getDebugStream(state) << "Found disableForkConditions for pc " << hexval(bb_start_pc) << "\n";
			signal->connect(sigc::mem_fun(*this, &GenerationSearcher::onDisableForkTargetExecution));
		}
    }

}
void GenerationSearcher::onOnlyForkTargetExecution(S2EExecutionState *state, uint64_t pc) {
	state->enableForking();
	getDebugStream(state) << "enableForking pc " << hexval(pc) << "\n";
	if(m_configedOnlyForkPoints.size() > 0){
		m_onlyForkPoints.insert(m_configedOnlyForkPoints.begin(),m_configedOnlyForkPoints.end());
	}
}

void GenerationSearcher::onDisableForkTargetExecution(S2EExecutionState *state, uint64_t pc) {
	state->disableForking();
	getDebugStream(state) << "disableForking pc " << hexval(pc) << "\n";
	m_onlyForkPoints.clear();
}

void GenerationSearcher::switchToCUPA() {
    m_multiSearcher->selectSearcher("CUPASearcher");
}

void GenerationSearcher::switchToGenerationSearcher() {
    m_multiSearcher->selectSearcher("GenerationSearcher");
}

bool GenerationSearcher::solveConstraints(S2EExecutionState *state, const klee::ref<klee::Expr> condition_e, klee::Assignment &assignment) {
	getDebugStream(state) << "Condition: " << condition_e << '\n';

	klee::ConstraintManager tmpConstraints = state->constraints;
	tmpConstraints.addConstraint(condition_e);

	klee::ArrayVec symbObjects = state->symbolics;
	std::vector<std::vector<unsigned char>> concreteObjects;
	std::map<klee::ArrayPtr, std::set<int> > pathrelated;

	klee::Solver *solver = state->solver()->solver;
	klee::Query query(tmpConstraints,
			klee::ConstantExpr::alloc(0, klee::Expr::Bool));
	if (!solver->getInitialValues(query, symbObjects, concreteObjects)) {
		getWarningsStream() << "Could not get symbolic solution\n";
		return false;
	}

	std::set< klee::ref<klee::ReadExpr> > all_reads;
	std::set< klee::ref<klee::Expr> > cset = tmpConstraints.getConstraintSet();
	for (auto acit = cset.begin(), acend = cset.end(); acit != acend; acit++) {
		std::vector < klee::ref<klee::ReadExpr> > reads;
		findReads(*acit, false, reads);
		all_reads.insert(reads.begin(), reads.end());
	}
	s2e()->getExecutor()->scanRelated(all_reads, pathrelated);

	for (unsigned i = 0; i < symbObjects.size(); ++i) {
		auto &arr = symbObjects[i];
		std::map<klee::ArrayPtr, std::set<int> >::const_iterator lit =
				pathrelated.find(arr);
		for (unsigned s = 0; s < arr->getSize(); ++s) {
			klee::ref<klee::Expr> e = state->concolics->evaluate(arr, s);
			if (isa<klee::ConstantExpr>(e)) {
				uint8_t val = dyn_cast<klee::ConstantExpr>(e)->getZExtValue();
				if (lit != pathrelated.end()) {
					std::set<int>::const_iterator offsetit = (*lit).second.find(
							s);
					if (offsetit != (*lit).second.end()) {

					} else {
						concreteObjects[i][s] = val;
					}
				} else {
					concreteObjects[i][s] = val;
				}
			}
		}
		assignment.add(symbObjects[i], concreteObjects[i]);
	}
	return true;
}

void GenerationSearcher::generationBugCase(S2EExecutionState *state, klee::ref<klee::Expr> symbolicAddress, uint64_t currentAddress){
	if(!m_searchBug) return;

	typedef std::pair<std::string, std::vector<unsigned char>> VarValuePair;
    typedef std::vector<VarValuePair> ConcreteInputs;
    uint64_t unsat_count = 0;
    uint64_t sat_count = 0;
	uint64_t currentAddress_up = currentAddress;
	uint64_t currentAddress_down = currentAddress;
	uint64_t currentAddress_max = currentAddress;
	uint64_t currentAddress_min = currentAddress;
	// get the upper
	while(true){
		if(unsat_count>=1){
			currentAddress_up = currentAddress_max + 1;
			sat_count = 0;
		}else{
			currentAddress_up += pow(2,sat_count);
		}
		//over flow
		if(currentAddress_up < currentAddress_max)
			break;
		klee::ref<klee::Expr> nextconcreteAddress = klee::ConstantExpr::create(
				currentAddress_up, symbolicAddress->getWidth());
		klee::ref<klee::Expr> conditionnext = klee::EqExpr::create(
				nextconcreteAddress, symbolicAddress);
		klee::ref<klee::Expr> conditionnextnot = klee::NotExpr::create(
				conditionnext);
		klee::Query query(state->constraints, conditionnextnot);
		bool truth;
		bool res = state->solver()->solver->mustBeTrue(query, truth);
		if (!res || truth) {
			unsat_count ++;
			if(unsat_count >= 2){
				break;
			}
			continue;
		}
		unsat_count = 0;
		sat_count ++;
		currentAddress_max = currentAddress_up;
	}

	if (currentAddress_max > currentAddress) {
		ConcreteInputs inputs;
		klee::Assignment assignment;
		klee::ref<klee::Expr> nextconcreteAddress = klee::ConstantExpr::create(
				currentAddress_max, symbolicAddress->getWidth());
		klee::ref<klee::Expr> conditionnext = klee::EqExpr::create(
				nextconcreteAddress, symbolicAddress);

		if (solveConstraints(state, conditionnext, assignment)) {
			for (const auto &it : assignment.bindings) {
				auto &array = it.first;
				const std::vector<unsigned char> &varData = it.second;
				inputs.push_back(std::make_pair(array->getName(), varData));
			}
			std::stringstream ss;
			ss << "testcase-" << "checkbug" << "-" << state->getID() << "-"
					<< hexval(state->regs()->getPc()) << "-"
					<< hexval(currentAddress_max);
			std::vector < std::string > fileNames;
			m_tc->assembleTestCaseToFiles(inputs, m_tc->getTemplates(state),
					ss.str(), fileNames);
			for (const auto &it : fileNames) {
				getDebugStream(state) << "Generated " << it << "\n";
			}
		}
	}
	unsat_count = 0;
	sat_count = 0;
	// get the min
	while(true){
		if(unsat_count>=1){
			currentAddress_down = currentAddress_min - 1;
			sat_count = 0;
		}else{
			currentAddress_down -= pow(2,sat_count);
		}
		//over flow
		if(currentAddress_down > currentAddress_min)
			break;
		klee::ref<klee::Expr> nextconcreteAddress = klee::ConstantExpr::create(
				currentAddress_down, symbolicAddress->getWidth());
		klee::ref<klee::Expr> conditionnext = klee::EqExpr::create(
				nextconcreteAddress, symbolicAddress);
		klee::ref<klee::Expr> conditionnextnot = klee::NotExpr::create(
				conditionnext);
		klee::Query query(state->constraints, conditionnextnot);
		bool truth;
		bool res = state->solver()->solver->mustBeTrue(query, truth);
		if (!res || truth) {
			unsat_count++;
			if (unsat_count >= 2) {
				break;
			}
			continue;
		}
		unsat_count = 0;
		sat_count++;
		currentAddress_min = currentAddress_down;
	}

	if (currentAddress_min < currentAddress) {
		ConcreteInputs inputs;
		klee::Assignment assignment;
		klee::ref<klee::Expr> nextconcreteAddress =
				klee::ConstantExpr::create(currentAddress_min,
						symbolicAddress->getWidth());
		klee::ref<klee::Expr> conditionnext = klee::EqExpr::create(
				nextconcreteAddress, symbolicAddress);
		if (solveConstraints(state, conditionnext, assignment)) {
			for (const auto &it : assignment.bindings) {
				auto &array = it.first;
				const std::vector<unsigned char> &varData = it.second;
				inputs.push_back(std::make_pair(array->getName(), varData));
			}
			std::stringstream ss;
			ss << "testcase-" << "checkbug" << "-" << state->getID() << "-"
					<< hexval(state->regs()->getPc()) << "-"
					<< hexval(currentAddress_min);
			std::vector < std::string > fileNames;
			m_tc->assembleTestCaseToFiles(inputs, m_tc->getTemplates(state),
					ss.str(), fileNames);
			for (const auto &it : fileNames) {
				getDebugStream(state) << "Generated " << it << "\n";
			}
		}
	}
}

void GenerationSearcher::onSymbolicAddress(S2EExecutionState *state, klee::ref<klee::Expr> virtualAddress, uint64_t concreteAddress,
                               bool &concretize, CorePlugin::symbolicAddressReason reason) {

	if (reason == CorePlugin::symbolicAddressReason::MEMORY) {
		 concretize = true;
		 bool check_bug = false;
		 auto moduleall = m_detector->getAllDescriptor(state);
		 auto module = m_detector->getCurrentDescriptor(state);
		 uint64_t nativePc;
		 if (moduleall) {
			 moduleall->ToNativeBase(state->regs()->getPc(), nativePc);
		 }

		 TranslationBlock *tb = state->getTb();
		 if(tb && module){
			 if(tb->se_tb_type == TB_JMP_IND){
				bool should_fork = false;
				m_currentForkReocord.push_back(state->regs()->getPc());
				if(m_forkHistory.size() > 0){
					m_forkHistory.pop_front();//erase(m_forkHistory.begin());
				}else{
					should_fork = true;
				}

				 uint64_t bb_start_pc = ((uint64_t) -1);
				 auto module = m_detector->getCurrentDescriptor(state);
				 if(module && module->ToNativeBaseRebased(tb->pc, bb_start_pc)){
					 s2e()->getInfoStream(state)
						<< "Forkable TB_JMP_IND while access SymbolicAddress at fps_bb_start_pc="
						<< hexval(bb_start_pc) << "\n";
					foreach2(rgit, m_forkRanges.begin(), m_forkRanges.end())
					{
						if (bb_start_pc == (*rgit).start) {
							should_fork = true;
							break;
						}
					}
				 }

				if(m_onlyForkConditions.size()>0){
					should_fork = false;
				}
				if (state->isForkingEnabled()) {
					should_fork = true;
				}

				if (should_fork) {
					concretize = false;
					state->m_forkAllSymbolicAddressAtOnce = true;
				}
			 }else{
				 //check bug
				 s2e()->getWarningsStream(state) << "Found symbolic memory access at pc=" << hexval(state->regs()->getPc()) << " " <<(moduleall ? moduleall->Name : "" )<< " nativePc=" << hexval(nativePc) << "\n";
				 check_bug = true;
			 }
		 }else{
			 //check bug
			 s2e()->getWarningsStream(state) << "Found symbolic memory access tb=null or module=nul at pc=" << hexval(state->regs()->getPc()) << "\n";
			 check_bug = true;
		 }
		 if(check_bug){
			 generationBugCase(state,virtualAddress,concreteAddress);
		 }
        return;
    }else{
        if (reason == CorePlugin::symbolicAddressReason::PC) {
            return;
        }
        klee::ref<klee::Expr> nextEip = virtualAddress;
        if (isa<klee::ConstantExpr>(nextEip)) {
            return;
        }
        concretize = true;
        //check bug
        s2e()->getWarningsStream(state) << "Found symbolic EIP at pc=" << hexval(state->regs()->getPc()) << "\n";
        generationBugCase(state,virtualAddress,concreteAddress);
    }
}

void GenerationSearcher::generateOffsetInfo(ExprList elist, const std::string &name_id) {
	std::string path;
	std::stringstream fileName;
	fileName << "offset_info-kill-" <<  name_id.c_str()  << ".json";
    path = s2e()->getOutputFilename(fileName.str());

	std::map<klee::ArrayPtr, std::set<int> > pathrelated;
	std::set<klee::ref<klee::ReadExpr> > all_reads;
	for (auto acit = elist.begin(), acend = elist.end(); acit != acend;
			acit++) {
		std::vector<klee::ref<klee::ReadExpr> > reads;
		findReads(*acit, false, reads);
		all_reads.insert(reads.begin(), reads.end());
	}
	s2e()->getExecutor()->scanRelated(all_reads, pathrelated);
	std::stringstream offset_str;
	QDict *pt = qdict_new();
	for (auto lit = pathrelated.begin(), litend = pathrelated.end();
			lit != litend; lit++) {
		QList *data = qlist_new();
		for (auto offset : (*lit).second) {
			qlist_append_obj(data, QOBJECT(qint_from_int(offset)));
		}
		qdict_put_obj(pt, (*lit).first->getName().c_str(), QOBJECT(data));
	}
	QString *json = qobject_to_json(QOBJECT(pt));
	offset_str << qstring_get_str(json) << "\n";
	QDECREF(json);
	QDECREF(pt);
	std::ofstream o(path.c_str());
	o << offset_str.str();
	o.close();
}

void GenerationSearcher::onUnableFork(S2EExecutionState *state, const klee::ref<klee::Expr> &condition){
	//generate testcase for current condition

	auto scondition = state->simplifyExpr(condition);
	auto module = m_detector->getCurrentDescriptor(state);
	if ( module){
		TranslationBlock *tb = state->getTb();
		bool allowFork = false;
		if (tb){
			uint64_t bb_start_pc = ((uint64_t) -1);
			if(module->ToNativeBaseRebased(tb->pc, bb_start_pc)){
				foreach2(rgit, m_forkRanges.begin(), m_forkRanges.end())
				{
					if (bb_start_pc == (*rgit).start) {
						allowFork = true;
						break;
					}
				}
			}
			if(m_onlyForkConditions.size()>0){
				allowFork = false;
			}
			if (state->isForkingEnabled()) {
				allowFork = true;
			}
		}
		//s2e()->getWarningsStream(state) << "Found onUnableFork at pc=" << hexval(state->regs()->getPc()) << "\n" << scondition << "\n";
		++m_forkCount[module->Name][state->regs()->getPc()];
		if (!allowFork && m_forkCount[module->Name][state->regs()->getPc()] > m_maxUnForkGen) {
			return;
		}
		ExprList elist;
		llvm::sys::TimeValue curTime = llvm::sys::TimeValue::now();
		elist.push_back(scondition);
		std::stringstream ss;
		ss << "notfork" << hexval(state->regs()->getPc()) << "|" << curTime.seconds();
		generateTestCaseForExprList(state, elist, ss.str());
		generateOffsetInfo(elist, ss.str());
		generateNotForkJsonHistoryFile(state, ss.str());
	}
}

void GenerationSearcher::onStateForkDecide(S2EExecutionState *state, bool *doFork) {
    auto module = m_detector->getCurrentDescriptor(state);
	auto moduleall = m_detector->getAllDescriptor(state);
	if (moduleall) {
		size_t pos = moduleall->Name.find(std::string("s2ecmd"));
		if (pos != std::string::npos) {
			return;
		}
	}

    if (!module) {
    	if(m_underConstraint)state->underConstraint = true;
    	*doFork = false;
        return;
    }

    if(state->m_forkAllSymbolicAddressAtOnce){
		return;
	}

    bool allowFork = false;
	m_currentForkReocord.push_back(state->regs()->getPc());
	if(m_forkHistory.size() > 0){
		m_forkHistory.pop_front();
    	allowFork = false;
    }else{
    	allowFork = true;
    }

    uint64_t curPc;
    if (!module->ToNativeBase(state->regs()->getPc(), curPc)) {
    	getDebugStream(state) << "Could not get relative pc for module " << module->Name << "\n";
    	if(m_underConstraint)state->underConstraint = true;
    	*doFork = false;
        return;
    }
    TranslationBlock *tb = state->getTb();
	int flags = 0; // 32-bit code by default

	uint64_t bb_start_pc = ((uint64_t) -1);
	if (tb){
		if(module->ToNativeBaseRebased(tb->pc, bb_start_pc)){
			foreach2(rgit, m_forkRanges.begin(), m_forkRanges.end())
			{
				if (bb_start_pc == (*rgit).start) {
					allowFork = true;
					break;
				}
			}
		}
		switch (state->getPointerSize()) {
		case 4:
			flags = 0;
			break;
		case 8:
			flags = 2;
			break;
		default:
			flags = 0;//pabort("Not supported code");
		}
		s2e()->getInfoStream(state) << "Forkable at fps_bb_start_pc=" << hexval(bb_start_pc)<< "\n";
	}
	if (m_onlyForkConditions.size() > 0) {
		allowFork = false;
	}

	if (state->isForkingEnabled() && m_onlyForkPoints.size() > 0) {
		if (state->isForkingEnabled() && m_onlyForkPoints.count(bb_start_pc)) {
			m_onlyForkPoints.erase(bb_start_pc);
			allowFork = true;
		} else {
			allowFork = false;
		}
	}
	csh handle;
	cs_mode mode;
	switch (flags) {
	case 0:
		mode = CS_MODE_32;
		break;
	case 1:
		mode = CS_MODE_16;
		break;
	case 2:
		mode = CS_MODE_64;
		break;
	default:
		mode = CS_MODE_32;
		break;
	}
	if(allowFork){
		if (m_unForkInstCache.find(curPc)!=m_unForkInstCache.end()){
			getDebugStream(state) << "Fork disabled for instruction by cache in module " << module->Name  << " at "<< hexval(curPc) << "\n";
			allowFork = false;
		}else{
			if (cs_open(CS_ARCH_X86, mode, &handle) == CS_ERR_OK) {
				cs_insn *insn;
				uint64_t pcaddress = state->regs()->getPc();
				size_t count;
				unsigned char code[15];
				size_t codesize = 15;
				if (state->mem()->read(pcaddress, code, codesize)) {
					count = cs_disasm(handle, code, codesize, pcaddress, 0, &insn);
					if (count) {
						if( insn[0].mnemonic[0]=='f'
							|| strcmp(insn[0].mnemonic,"movapd")==0
							|| strcmp(insn[0].mnemonic,"movupd")==0
							|| strcmp(insn[0].mnemonic,"movlpd")==0
							|| strcmp(insn[0].mnemonic,"movhpd")==0
							|| strcmp(insn[0].mnemonic,"addpd")==0
							|| strcmp(insn[0].mnemonic,"addsd")==0
							|| strcmp(insn[0].mnemonic,"subpd")==0
							|| strcmp(insn[0].mnemonic,"subsd")==0
							|| strcmp(insn[0].mnemonic,"mulpd")==0
							|| strcmp(insn[0].mnemonic,"mulsd")==0
							|| strcmp(insn[0].mnemonic,"divpd")==0
							|| strcmp(insn[0].mnemonic,"divsd")==0
							|| strcmp(insn[0].mnemonic,"sqrtpd")==0
							|| strcmp(insn[0].mnemonic,"sqrtsd")==0
							|| strcmp(insn[0].mnemonic,"maxpd")==0
							|| strcmp(insn[0].mnemonic,"maxsd")==0
							|| strcmp(insn[0].mnemonic,"minpd")==0
							|| strcmp(insn[0].mnemonic,"minsd")==0
							|| strcmp(insn[0].mnemonic,"andpd")==0
							|| strcmp(insn[0].mnemonic,"andnpd")==0
							|| strcmp(insn[0].mnemonic,"orpd")==0
							|| strcmp(insn[0].mnemonic,"xorpd")==0
							|| strcmp(insn[0].mnemonic,"cvtps2pd")==0
							|| strcmp(insn[0].mnemonic,"cvtss2sd")==0
							|| strcmp(insn[0].mnemonic,"cvtpd2ps")==0
							|| strcmp(insn[0].mnemonic,"cvtsd2ss")==0
							|| strcmp(insn[0].mnemonic,"cvtpd2pi")==0
							|| strcmp(insn[0].mnemonic,"cvtpi2pd")==0
							|| strcmp(insn[0].mnemonic,"cvtpd2dq")==0
							|| strcmp(insn[0].mnemonic,"cvtdq2pd")==0
							|| strcmp(insn[0].mnemonic,"cvtsd2si")==0
							|| strcmp(insn[0].mnemonic,"cvtsi2sd")==0
							|| strcmp(insn[0].mnemonic,"cvtps2dq")==0
							|| strcmp(insn[0].mnemonic,"cvtdq2ps")==0
							|| strcmp(insn[0].mnemonic,"ucomisd")==0
						){
							m_unForkInstCache.insert(curPc);
							getDebugStream(state) << "Fork disabled for instruction " << insn[0].mnemonic << " in module " << module->Name  << " at "<< hexval(curPc) << "\n";
							allowFork = false;
						}
						cs_free(insn, count);
					} else {
					}
				}
				cs_close(&handle);
			}
		}
	}
	if(allowFork){
		s2e()->getDebugStream(state)<< "Allow StateForkDecide Native pc= " << hexval(curPc)<< "\n";
	}else{
		s2e()->getDebugStream(state) << "Disable StateForkDecide Native pc= " << hexval(curPc)<< "\n";
		if(m_underConstraint)state->underConstraint = true;
		*doFork = false;
	}
}

void GenerationSearcher::onStateFork(S2EExecutionState *oldState, const std::vector<S2EExecutionState *> &newStates,
                               const std::vector<klee::ref<klee::Expr>> &newConditions) {
    if (oldState == m_initialState) {
    	uint64_t hitpoint = oldState->regs()->getPc();
    	int count = 0;
        for (auto it : newStates) {
            if (it != oldState) {
				it->m_should_del = true;
				if(!oldState->m_forkAllSymbolicAddressAtOnce)
				{
					if (multiHitConstraints.find(hitpoint) != multiHitConstraints.end()) {
						multiHitConstraints[hitpoint].push_back(newConditions[count]);
						multiHitRecord[hitpoint] = multiHitRecord[hitpoint] + 1;
					}else{
						ExprList elist;
						elist.push_back(newConditions[count]);
						multiHitConstraints[hitpoint] = elist;
						multiHitRecord[hitpoint] = 1;
					}
				}

				if (!it->isZombie()) {
					m_tc->generateTestCases(it, "kill", testcases::TC_FILE);
					generateJsonHistoryFile(it);
					try {
						s2e()->getExecutor()->terminateState(*it,
								"terminate write generation case.");
					} catch (s2e::CpuExitException&) {
					} catch (...) {
					}
				}
            }
            count ++;
        }
    }
}

std::string GenerationSearcher::generateJsonHistoryFile(S2EExecutionState *state) {
    std::string path;
    std::stringstream fileName;
    fileName << "forkhistory-kill-" << state->getID() << ".json";
    path = s2e()->getOutputFilename(fileName.str());

    generateJsonHistoryFile(state, path);

    return path;
}

std::string GenerationSearcher::generateNotForkJsonHistoryFile(S2EExecutionState *state,const std::string &name_id) {
    std::string path;
    std::stringstream fileName;
    fileName << "forkhistory-kill-" << name_id.c_str() << ".json";
    path = s2e()->getOutputFilename(fileName.str());

    generateJsonHistoryFile(state, path);

    return path;
}

void GenerationSearcher::generateJsonHistoryFile(S2EExecutionState *state, const std::string &path) {
    std::stringstream history;
    generateJsonHistory(state, history);

    std::ofstream o(path.c_str());
    o << history.str();
    o.close();
}

void GenerationSearcher::generateJsonHistory(S2EExecutionState *state, std::stringstream &history) {
    QDict *pt = qdict_new();

    QList *forkPoints = qlist_new();
    for (auto forkPoint : m_currentForkReocord) {
        qlist_append_obj(forkPoints, QOBJECT(qint_from_int(forkPoint)));
    }
    qdict_put_obj(pt, "forkPoints", QOBJECT(forkPoints));

    QString *json = qobject_to_json(QOBJECT(pt));

    history << qstring_get_str(json) << "\n";

    QDECREF(json);
    QDECREF(pt);
}

void GenerationSearcher::generateTestCaseForExprList(S2EExecutionState *state, ExprList explist, const std::string &prefix) {
	klee::ArrayVec symbObjects = state->symbolics;
	klee::ConstraintManager oldConstraints = state->constraints;
	std::map<klee::ArrayPtr, std::set<int> > pathrelated;
	klee::ConstraintManager constraints;
	std::set<klee::ref<klee::ReadExpr> > all_reads;
	foreach2 (ite, explist.begin(), explist.end())
	{
		constraints.addConstraint(*ite);
		std::vector<klee::ref<klee::ReadExpr> > reads;
		findReads(*ite, false, reads);
		all_reads.insert(reads.begin(), reads.end());
	}
	s2e()->getExecutor()->scanRelated(all_reads, pathrelated);
	state->constraints = constraints;
	std::vector < std::vector<unsigned char> > concreteObjects;
	klee::Assignment assignment;
	ConcreteInputs inputs;
	auto solver = state->solver()->solver;
	if (solver->getInitialValues(
			klee::Query(constraints,
					klee::ConstantExpr::alloc(0, klee::Expr::Bool)),
			symbObjects, concreteObjects)) {
		for (unsigned i = 0; i < symbObjects.size(); ++i) {
			auto &arr = symbObjects[i];
			std::map<klee::ArrayPtr, std::set<int> >::const_iterator lit =
					pathrelated.find(arr);
			for (unsigned s = 0; s < arr->getSize(); ++s) {
				klee::ref<klee::Expr> e = state->concolics->evaluate(arr, s);
				if (isa<klee::ConstantExpr>(e)) {
					uint8_t val =
							dyn_cast<klee::ConstantExpr>(e)->getZExtValue();
					if (lit != pathrelated.end()) {
						std::set<int>::const_iterator offsetit =
								(*lit).second.find(s);
						if (offsetit != (*lit).second.end()) {
						} else {
							concreteObjects[i][s] = val;
						}
					} else {
						concreteObjects[i][s] = val;
					}
				}
			}
			assignment.add(symbObjects[i], concreteObjects[i]);
		}
		for (const auto &it : assignment.bindings) {
			auto &array = it.first;
			const std::vector<unsigned char> &varData = it.second;
			inputs.push_back(std::make_pair(array->getName(), varData));
		}
		std::stringstream ss;
		ss << "testcase-" << prefix << "-" << state->getID();
		std::vector < std::string > fileNames;
		m_tc->assembleTestCaseToFiles(inputs, m_tc->getTemplates(state),
				ss.str(), fileNames);
		for (const auto &it : fileNames) {
			getDebugStream(state) << "Generated " << it << "\n";
		}
	}
	state->constraints = oldConstraints;
}

void GenerationSearcher::onStateKill(S2EExecutionState *state) {
	if (state == m_initialState) {
		try {
			state->m_should_del = true;
			if(m_genMultiHitCase){
				for (auto &it : multiHitRecord) {
					const auto &hitpoint = it.first;
					const auto &hitcount = it.second;
					if (hitcount > 3) {
						ExprList explist = multiHitConstraints[hitpoint];
						std::stringstream ss;
						ss << "hitpoint" << hexval(hitpoint);
						generateTestCaseForExprList(state, explist, ss.str());
					}
				}
			}
		} catch (...) {

		}
		m_initialState = nullptr;
	}
}

void GenerationSearcher::update(klee::ExecutionState *current, const klee::StateSet &addedStates,
                          const klee::StateSet &removedStates) {
    S2EExecutionState *cs = dynamic_cast<S2EExecutionState *>(current);
    if (!m_initialState) {
        if (cs->getID() == 0) {
            m_initialState = cs;
        	if(m_onlyForkConditions.size()>0){
        		m_initialState->disableForking();
        	}
        }
    }

    for (auto removedState : removedStates) {
        S2EExecutionState *es = dynamic_cast<S2EExecutionState *>(removedState);
        if (es == m_initialState) {
			m_initialState = nullptr;
		}
    }
}

klee::ExecutionState &GenerationSearcher::selectState() {
     return *m_initialState;
}

bool GenerationSearcher::empty() {
    return m_initialState == nullptr || m_initialState->m_should_del;
}

} // namespace generations
} // namespace plugins
} // namespace s2e
